from datasketch import MinHash, MinHashLSH
import redis
import pickle as pkl
from typing import List, Dict
import logging
from tqdm import tqdm

class MinHashRedis:
    def __init__(self, 
                 num_perm: int = 128,
                 threshold: float = 0.5,
                 redis_host: str = 'localhost',
                 redis_port: int = 6379,
                 password: str = '',
                 redis_db: int = 0,
                 expire_time: int = 86400 * 30,
                 basename: str = 'mylsh'
                 ):
        
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)
        
        self.redis_raw = redis.Redis(
            host=redis_host,
            port=redis_port,
            password=password,
            db=redis_db
        )
        
        self.expire_time = expire_time
        self.num_perm = num_perm
        self.threshold = threshold
        self.basename = basename.encode()
        
        self.lsh = MinHashLSH(
            threshold=self.threshold,
            num_perm=self.num_perm,
            storage_config={
                'type': 'redis',
                'basename': self.basename,
                'redis': {
                    'host': redis_host,
                    'port': redis_port,
                    'password': password,
                    'db': redis_db,
                }
            }
        )
        
    def _create_minhash(self, items: List[str]) -> MinHash:
        minhash = MinHash(num_perm=self.num_perm)
        for item in items:
            minhash.update(item.encode('utf8'))
        return minhash
    
    def add_user(self, user_id: str, apps: List[str], force: bool = False) -> bool:
        if self.check_user_exists(user_id):
            self.logger.info(f"用户 {user_id} 已存在")
            if not force:
                self.logger.info(f"用户 {user_id} 跳过添加")
                return True
            self.logger.info(f"用户 {user_id} 强制更新")
            minhash = self._create_minhash(apps)
            minhash_key = f"minhash:{self.basename.decode()}:{user_id}"
            self.lsh.remove(minhash_key)
        else:
            self.logger.info(f"用户 {user_id} 已添加")
            minhash = self._create_minhash(apps)
            minhash_key = f"minhash:{self.basename.decode()}:{user_id}"
        self.redis_raw.set(minhash_key, pkl.dumps(minhash))
        self.redis_raw.expire(minhash_key, self.expire_time)
        try:    
            self.lsh.insert(minhash_key, minhash)
        except Exception as e:
            self.logger.error(f"插入minhash失败: {str(e)}")
            return False
        return True
    
    def check_user_exists(self, user_id: str) -> bool:
        try:
            minhash_key = f"minhash:{self.basename.decode()}:{user_id}"
            return self.redis_raw.exists(minhash_key)
        except Exception as e:
            self.logger.error(f"检查用户存在状态失败: {str(e)}")
            return False
    
    def find_similar_users(self, 
                            apps: List[str], 
                            min_similarity: float = 0.0) -> List[Dict]:
            query_minhash = self._create_minhash(apps)
            candidates = self.lsh.query(query_minhash)
            
            with self.redis_raw.pipeline() as pipe:
                for candidate_id in candidates:
                    pipe.get(candidate_id)
                minhash_data_list = pipe.execute()
            
            results = []
            for candidate_id, minhash_data in zip(candidates, minhash_data_list):
                if not minhash_data:
                    continue
                
                candidate_minhash = pkl.loads(minhash_data)
                similarity = query_minhash.jaccard(candidate_minhash)
                if similarity >= min_similarity:
                    results.append({
                        'minhash_key': candidate_id,
                        'similarity': similarity
                    })
                    
            return sorted(results, key=lambda x: x['similarity'], reverse=True)
            
    
    def find_and_add(self, user_id: str, apps: List[str], min_similarity: float = 0.0, force=False):
        similar_users = self.find_similar_users(apps, min_similarity)
        self.add_user(user_id, apps, force=force)
        return similar_users